# _*_ coding: utf-8 _*_
# 导入所需的库
import os
import random
import sys

# 导入数据处理和可视化库
import matplotlib.pyplot as plt
import numpy as np

# 导入深度学习框架 PyTorch 相关库
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from ResNet import resnet18
# from ResNet import resnet34
# from ResNet import resnet50
# import torchvision.models as models
from tqdm import tqdm

import seaborn as sns
from sklearn.metrics import confusion_matrix

# 设置随机种子以保证结果的可重复性
def setup_seed(seed):
    np.random.seed(seed)  # 设置 Numpy 随机种子
    random.seed(seed)  # 设置 Python 内置随机种子
    os.environ['PYTHONHASHSEED'] = str(seed)  # 设置 Python 哈希种子
    torch.manual_seed(seed)  # 设置 PyTorch 随机种子
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)  # 设置 CUDA 随机种子
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.benchmark = False  # 关闭 cuDNN 的自动优化模式，适用于输入尺寸不固定的场景。如果你的输入尺寸固定，开启 benchmark 可以提高性能。
        # benchmark=True，cuDNN 会在每次运行时自动寻找最适合当前输入尺寸和硬件配置的卷积算法。这种优化可以提高运行速度，但会增加一些初始开销。
        # 当 benchmark=False 时，cuDNN 不会进行这种自动优化，而是使用默认的算法。
        torch.backends.cudnn.deterministic = True  # 启用 cuDNN 的确定性算法模式，确保结果可重复。


# 设置随机种子
setup_seed(0)
# 检查是否有可用的 GPU，如果有则使用 GPU，否则使用 CPU
if torch.cuda.is_available():
    device = torch.device("cuda")  # 使用 GPU
    print("CUDA is available. Using GPU.")
else:
    device = torch.device("cpu")  # 使用 CPU
    print("CUDA is not available. Using CPU.")

transform = {
    "train_model1": transforms.Compose([transforms.RandomResizedCrop(size=224),  # 随机裁剪和缩放
                                 transforms.RandomHorizontalFlip(),  # 随机水平翻转
                                 # transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 随机颜色变换,亮度、对比度、饱和度和色调
                                 # transforms.RandomRotation(15),  # 随机旋转
                                 transforms.ToTensor(),  # 转换为张量
                                 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
    "valid": transforms.Compose([transforms.Resize((224, 224)),
                                transforms.ToTensor(),
                                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
}

train_dataset = datasets.ImageFolder("./dataset_peach/train_model1", transform=transform["train_model1"])
valid_dataset = datasets.ImageFolder("./dataset_peach/valid", transform=transform["valid"])

train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)
valid_dataloader = DataLoader(valid_dataset, batch_size=4, shuffle=False)

# 打印一下图片
examples = enumerate(valid_dataloader)  # 将 valid_dataloader 转换为一个枚举对象，方便获取批次的索引和数据。
batch_idx, (imgs, labels) = next(examples)  # 获取验证集的第一个批次数据。
# batch_idx=0，当前批次的索引（从 0 开始）。
# imgs：当前批次的图像数据，形状为 ( channels, height, width)。

# print(f"batch_idx: {batch_idx},(imgs, labels):{(imgs, labels)}")
#--------------------------------------------------------------------------------------------
# for i in range(4):
#     # 反归一化并显示图像
#     mean = np.array([0.485, 0.456, 0.406])
#     std = np.array([0.229, 0.224, 0.225])
#     # 归一化后的值再反归一化，得到图像初始数据，std[:, None, None] 和 mean[:, None, None]：将均值和标准差扩展为与图像相同的形状
#     image = imgs[i].numpy() * std[:, None, None] + mean[:, None, None]
#     # 将图片转成numpy数组，image形状为 (channels, height, width)，转换为 (height, width, channels)，以便 matplotlib 可以正确显示。
#     image = np.transpose(image, (1, 2, 0))
#     # 将图像像素值从 [0, 1] 范围转换为 [0, 255] 范围，并转换为整数类型。
#     image = (image * 255).astype(np.uint8)
#     plt.subplot(2, 2, i+1)
#     plt.imshow(image)
#     plt.title(f"Truth: {labels[i]}")
# plt.show()
#--------------------------------------------------------------------------------------------

# //////////这段代码加载了一个预训练的ResNet18模型，并修改了其最后的全连接层，以适应一个只有3个类的新的分类任务。这种修改是迁移学习中的一种常见做法，
# //////////即利用在大型数据集上预训练的模型的特征提取能力，并通过调整最后几层来适应新的任务。
model = resnet18().to(device)
# model = resnet34().to(device)
# model = resnet50().to(device)
# model = models.resnet18(pretrained=True).to(device)  # 使用ResNet-18

# 加载一个预训练的模型权重文件，并将这些权重应用到模型上，以便使用预训练的模型进行后续的任务，如图像分类、特征提取等。
# 将加载的权重应用到模型上
model.load_state_dict(torch.load('./pre_model/resnet18-f37072fd.pth', weights_only=False))
# model.load_state_dict(torch.load('./pre_model/resnet34-333f7ec4.pth'))
# model.load_state_dict(torch.load('./pre_model/resnet50-19c8e357.pth'))

# 获取原始ResNet18模型中最后全连接层的输入特征数，在ResNet18中，这个全连接层通常用于将卷积层提取的特征映射到1000个ImageNet分类上。
# 但是，对于新的分类任务，我们可能需要不同数量的输出类。
# 好处是可以利用预训练模型的特征提取能力，同时避免重新训练整个模型。
# 冻结预训练模型的大部分层，仅训练最后几层（通常是FC层）以适应新任务。通过自定义FC层，我们可以轻松地实现这一点，而无需重新创建整个模型。












# 获取 model 模型中 fc 全连接层的输入特征数量，并将其赋值给变量 fc_inputs
fc_inputs = model.fc.in_features
# 用一个新的全连接层替换了原始的全连接层，使用nn.Sequential来包装新的全连接层。
# model.fc = nn.Sequential(
#     # 将输入特征（fc_inputs）映射到3个输出类上
#     nn.Linear(fc_inputs, 3),
# ).to(device)


model.fc = nn.Sequential(
    nn.Linear(fc_inputs, 128),  # 增加一个隐藏层
    nn.ReLU(),  # 添加 ReLU 激活函数
    # nn.Dropout(0.5),  # 添加 Dropout 防止过拟合
    nn.Linear(128, 3),  # 最终映射到 3 个输出类上
    nn.Softmax(dim=1)  # 分类任务使用 Softmax 输出概率分布
).to(device)


# //////////
# //////////
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)

train_num = len(train_dataset)
valid_num = len(valid_dataset)

epochs = 10
for epoch in range(epochs):
    most_acc = 0.0
    # 训练
    model.train()
    train_acc = 0.0
    train_loss = 0.0
    # 进度条功能，train_bar将 train_dataloader 包装为一个带有进度条的可迭代对象。
    train_bar = tqdm(train_dataloader, file=sys.stdout)
    # 生成epoch[1/10] 的字符串
    spaces = " " * (len(f"epoch[{epoch + 1}/{epochs}]") + 1)
    # step表示当前迭代的步数从0开始计数），data表示从train_bar封装的train_dataloader中取出的一个批次的数据，同时更新进度条。
    for step, data in enumerate(train_bar):
        images, labels = data
        # 将数据移动到设备上
        images = images.to(device)
        labels = labels.to(device)

        train_out = model(images)
        # print(train_out.shape)
        loss = criterion(train_out, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # loss：打印整个张量的信息，包括张量的值、形状、数据类型和设备信息；loss.item()打印张量的标量值（仅适用于包含单个值的张量）
        train_loss += loss.item()

        # train_out:Tensor(16, 3), dim=1,计算每行的最大值。每行是三分类的概率，获取概率最大的分类结果，返回张量的最大值以及对应的索引
        predict_y = torch.max(train_out, dim=1)[1]

        # predict_y和labels对应位置的元素如果相等则返回1，不相等返回0。求和也就是求相等的个数
        train_acc += torch.eq(predict_y, labels).sum().item()

        train_bar.desc = f"epoch[{epoch + 1}/{epochs}] train_loss:{loss:.3f}"
    # print(train_acc)
    # print(train_num)
    train_loss /= train_num
    train_acc /= train_num
    print(f'{spaces}train_loss: {train_loss:.3f}  train_acc: {train_acc:.3f}')

    # 验证
    model.eval()
    val_acc = 0.0  # accumulate accurate number / epoch
    val_loss = 0.0
    with torch.no_grad():
        val_bar = tqdm(valid_dataloader, file=sys.stdout)
        for val_data in val_bar:
            val_images, val_labels = val_data
            # 将数据移动到设备上
            val_images = val_images.to(device)
            val_labels = val_labels.to(device)

            val_out = model(val_images)
            loss = criterion(val_out, val_labels)
            val_loss += loss.item()

            predict_y = torch.max(val_out, dim=1)[1]
            val_acc += torch.eq(predict_y, val_labels).sum().item()

            val_bar.desc = f"{spaces}val_loss:{loss:.3f}"

    val_loss /= valid_num
    val_acc /= valid_num
    print(f'{spaces}val_loss: {val_loss:.3f}  val_acc: {val_acc:.3f}')

    if val_acc > most_acc:
        most_acc = val_acc
        # 手动创建model文件夹
        torch.save(model.state_dict(), './model/best.pth')
    torch.save(model.state_dict(), './model/last.pth')

# 训练完成后，加载最好的模型进行评估
model.load_state_dict(torch.load('./model/best.pth'))
# 评估模型
correct = 0
total = 0
predicted_labels = []
true_labels = []
model.eval()
with torch.no_grad():
    for images, labels in valid_dataloader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        # 将当前批次的预测标签predicted从设备（如GPU）移动到CPU上，并转换为NumPy数组，然后将其添加到列表predicted_labels中
        predicted_labels.extend(predicted.cpu().numpy())
        true_labels.extend(labels.cpu().numpy())

print('Accuracy of the model on the test images: {} %'.format(100 * correct / total))

# 生成混淆矩阵
conf_matrix = confusion_matrix(true_labels, predicted_labels)

# 可视化混淆矩阵
plt.figure(figsize=(10, 8))
sns.heatmap(conf_matrix, annot=True, fmt="d", cmap="Blues", cbar=False)
plt.xlabel("Predicted Labels")
plt.ylabel("True Labels")
plt.title("Confusion Matrix")
plt.show()

